#!/bin/bash

# User Plane Traffic Generation Script
# This script generates various types of user plane traffic for interface capture testing

set -e

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Function to print colored output
print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Function to check if VM4 is accessible
check_vm4_access() {
    if ! ssh -o ConnectTimeout=5 ubuntu@vm4 "echo 'OK'" > /dev/null 2>&1; then
        print_error "Cannot access VM4. Please check SSH connectivity."
        exit 1
    fi
}

# Function to start continuous ping
start_continuous_ping() {
    print_status "Starting continuous ping to 8.8.8.8..."
    print_status "Ping output will be visible in real-time..."
    
    # Start ping in background but show output
    ssh ubuntu@vm4 "ping -i 1 8.8.8.8" &
    local ping_pid=$!
    
    # Wait a moment to see initial output
    sleep 3
    
    print_success "Continuous ping started (PID: $ping_pid)"
    print_status "To stop ping, run: ssh ubuntu@vm4 'pkill ping'"
}

# Function to generate HTTP traffic
generate_http_traffic() {
    print_status "Generating HTTP traffic..."
    
    # Test HTTP connectivity
    if ssh ubuntu@vm4 "curl -s --connect-timeout 5 http://www.google.com > /dev/null"; then
        print_success "HTTP connectivity confirmed"
        
        # Start periodic HTTP requests
        ssh ubuntu@vm4 "while true; do curl -s --connect-timeout 5 http://www.google.com > /dev/null; sleep 2; done &"
        print_success "Periodic HTTP traffic started"
    else
        print_warning "HTTP connectivity failed, skipping HTTP traffic generation"
    fi
}

# Function to generate DNS traffic
generate_dns_traffic() {
    print_status "Generating DNS traffic..."
    
    # Test DNS resolution
    if ssh ubuntu@vm4 "nslookup google.com > /dev/null 2>&1"; then
        print_success "DNS resolution confirmed"
        
        # Start periodic DNS lookups
        ssh ubuntu@vm4 "while true; do nslookup google.com > /dev/null 2>&1; sleep 3; done &"
        print_success "Periodic DNS traffic started"
    else
        print_warning "DNS resolution failed, skipping DNS traffic generation"
    fi
}

# Function to generate multiple ping targets
generate_multiple_pings() {
    print_status "Generating multiple ping targets..."
    
    # Ping different targets
    ssh ubuntu@vm4 "while true; do ping -c 1 1.1.1.1 > /dev/null 2>&1; sleep 2; done &"
    ssh ubuntu@vm4 "while true; do ping -c 1 9.9.9.9 > /dev/null 2>&1; sleep 3; done &"
    ssh ubuntu@vm4 "while true; do ping -c 1 208.67.222.222 > /dev/null 2>&1; sleep 4; done &"
    
    print_success "Multiple ping targets started"
}

# Function to verify traffic
verify_traffic() {
    print_status "Verifying traffic generation..."
    
    # Check if ping is running
    if ssh ubuntu@vm4 "pgrep ping > /dev/null"; then
        print_success "Ping process is running"
    else
        print_warning "Ping process not found"
    fi
    
    # Check traffic on UPF
    print_status "Checking traffic on UPF (VM2)..."
    ssh ubuntu@vm2 "sudo tcpdump -i ogstun -c 5 -n" 2>/dev/null || print_warning "No traffic detected on UPF"
    
    # Check traffic on UE
    print_status "Checking traffic on UE (VM4)..."
    ssh ubuntu@vm4 "sudo tcpdump -i uesimtun0 -c 5 -n" 2>/dev/null || print_warning "No traffic detected on UE"
}

# Function to show traffic status
show_status() {
    print_status "Current traffic generation status:"
    
    echo "Background processes on VM4:"
    ssh ubuntu@vm4 "jobs -l" 2>/dev/null || echo "No background jobs found"
    
    echo ""
    echo "Network connectivity test:"
    ssh ubuntu@vm4 "ping -c 3 8.8.8.8" 2>/dev/null || print_warning "Ping test failed"
}

# Function to stop all traffic
stop_traffic() {
    print_status "Stopping all traffic generation..."
    
    # Stop all background processes
    ssh ubuntu@vm4 "pkill ping" 2>/dev/null || true
    ssh ubuntu@vm4 "pkill curl" 2>/dev/null || true
    ssh ubuntu@vm4 "pkill nslookup" 2>/dev/null || true
    
    # Kill all background jobs
    ssh ubuntu@vm4 "jobs -p | xargs -r kill" 2>/dev/null || true
    
    print_success "All traffic generation stopped"
}

# Main execution
main() {
    echo "=========================================="
    echo "User Plane Traffic Generation Script"
    echo "=========================================="
    
    # Check VM4 accessibility
    check_vm4_access
    
    # Start traffic generation
    start_continuous_ping
    generate_http_traffic
    generate_dns_traffic
    generate_multiple_pings
    
    # Verify traffic
    verify_traffic
    
    # Show status
    show_status
    
    print_success "Traffic generation setup completed!"
    echo ""
    print_status "Traffic is now being generated continuously."
    print_status "You can now run interface captures."
    print_status "To stop traffic, run: ./scripts/generate-traffic.sh --stop"
    print_status "To check status, run: ./scripts/generate-traffic.sh --status"
}

# Handle command line arguments
case "${1:-}" in
    --stop)
        stop_traffic
        ;;
    --status)
        show_status
        ;;
    --help|-h)
        echo "Usage: $0 [OPTION]"
        echo ""
        echo "Options:"
        echo "  --stop     Stop all traffic generation"
        echo "  --status   Show current traffic status"
        echo "  --help     Show this help message"
        echo ""
        echo "Default: Start traffic generation"
        ;;
    *)
        main "$@"
        ;;
esac 